import logging
import os.path
import sys
import multiprocessing
from gensim.corpora import WikiCorpus
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
import json

logger = logging.getLogger("train")
logging.basicConfig(format='%(asctime)s: %(levelname)s: %(message)s')
logging.root.setLevel(level=logging.INFO)
logger.info("running %s" % ' '.join(sys.argv))

jieba_cut_data_path = 'jieba_cut.txt'
word2vec_save_path = 'model.vec'


'''
size: word向量的维度。
window: 一个句子中当前单词和被预测单词的最大距离。
min_count: 忽略词频小于此值的单词。
workers:训练模型时使用的线程数。
sg: 模型的训练算法: 1: skip-gram; 0: CBOW.
hs: 1: 采用hierarchical softmax训练模型; 0: 使用负采样。
negative : 使用负采样，设置多个负采样(通常在5-20之间)。
cbow_mean: – 0: 使用上下文单词向量的总和; 1: 使用均值，适用于使用CBOW。
alpha : 初始学习率。
iter (int, optional) – 迭代次数。
sorted_vocab ({0, 1}, optional) – 如果为1，则在分配单词索引前按降序对词汇表进行排序。
'''

curpus_path = r'/home/qiufengfeng/nlp/competition/datagrand/基于大规模预训练模型的风险事件标签识别/datagrand_2021_unlabeled_data.json'

class DatagrandSentence():
    def __init__(self,corpus_path):
        self.corpus_path = corpus_path

    def __iter__(self):
        with open(self.corpus_path, encoding='utf-8') as fread:
            for line in fread:
                json_obj = json.loads(line.strip())
                text = '[CLS] ' + json_obj['title'] + " [SEP]"
                yield text.split()

model = Word2Vec(DatagrandSentence(curpus_path), size=768, window=5, min_count=1,workers=multiprocessing.cpu_count(),iter=30)

model.wv.save_word2vec_format(word2vec_save_path, binary=False)